% For a DSGE model approximated up to third order, this function calculates 
% the conditional covariances based on the the pruning
% scheme. The system reads
%     xf(t+1)  = hx*xf(t)  + sig*eta*eps(t);
%     xs(t+1)  = hx*xs(t)  + 0.5*reshape(hxx,nx,nx^2)*kron(xf(t),xf(t))+ 1/2*sig^2*hss;
%     xrd(t+1) = hx*xrd(t) + 2*0.5*reshape(hxx,nx,nx^2)*kron(xf(t),xs(t))...
%                + 1/6*reshape(hxxx,nx,nx^3)*kron(xf(t),kron(xf(t),xf(t)))
%                + 3/6*hssx*sig^2*xf(t) + 1/6*sig^3*hsss;
%     y(t+1)  = gx*(xf(t+1)+xs(t+1)+xrd(t+1)) ...
%                + 0.5*reshape(gxx,ny,nx^2)*(kron(xf(t+1),xf(t+1)+2*kron(xf(t+1),xs(t+1))+ 1/2*sig^2*gss
%                + 1/6*reshape(gxxx,ny,nx^3)*(kron(xf(t+1),kron(xf(t+1),xf(t+1)))
%                + 3/6*gssx*sig^2*xf(t+1) + 1/6*sig^3*gsss;
%
% IMPORTANT: Here we allow for general shock distributions, i.e. they may be
% non-symmetric. 
%
% INPUTS:
% gx,...,sig : The second order approximation of the DSGE model
% vectorMom3 : vector with dimensions 1*ne with third moments of shocks
% vectorMom4 : vector with dimensions 1*ne with fourth moments of shocks
% vectorMom5 : vector with dimensions 1*ne with fifth moments of shocks
% vectorMom6 : vector with dimensions 1*ne with sixth moments of shocks
% auto_lag   : the lag lenght for the auto-covariances (they need to be divided
%              by the standard deviations to get the correlations)
%
% OUTPUTS: 
% See below
%
% By M. Andreasen, 2019

function out = conditionalVariance_3rd(xSim,model,addBondYieldsON)

if nargin == 2
    addBondYieldsON = 1;
end
    
% The dimensions
sig     = 1;
eta     = model.eta;  
[nx,ne] = size(eta);
sigeta  = sig*eta;

% Adding bond yields to the g-function
if addBondYieldsON == 1
    gx    = [model.gx;model.byx(model.maturities,:)];
    nyOld = size(model.gx,1);
    ny    = size(gx,1);
    gxx   = zeros(ny,nx,nx);
    gxxx  = zeros(ny,nx,nx,nx);
    gxx(1:nyOld,:,:)   = model.gxx;
    gxxx(1:nyOld,:,:,:)= model.gxxx;
    for i=1:length(model.maturities)
        gxx(nyOld+i,:,:)    = reshape(model.byxx(model.maturities(1,i),:,:),nx,nx);
        gxxx(nyOld+i,:,:,:) = reshape(model.byxxx(model.maturities(1,i),:,:,:),nx,nx,nx);
    end
    gss  = [model.gss ;model.byss(model.maturities,:)];
    gssx = [model.gssx;model.byssx(model.maturities,:)];
    gsss = [model.gsss;model.bysss(model.maturities,:)];
else
    gx   = model.gx;
    gxx  = model.gxx;
    gxxx = model.gxxx;
    gss  = model.gss;
    gssx = model.gssx;
    gsss = model.gsss;
    ny    = size(gx,1);
end

hx   = model.hx;
hxx  = model.hxx;
hxxx = model.hxxx;
hss  = model.hss;
hssx = model.hssx;
hsss = model.hsss;
  
% The moments in the normal distribution
vectorMom3 = zeros(1,ne);
vectorMom4 = 3*ones(1,ne);
vectorMom5 = zeros(1,ne);
vectorMom6 = 15*ones(1,ne);

% The value of E_u
E_eps3 = zeros(ne^3,1);
index = 0;
for phi1=1:ne
    for phi2=1:ne
        for phi3=1:ne
            index = index + 1;
            if phi1 == phi2 && phi1 == phi3
                E_eps3(index,1) = vectorMom3(1,phi1);
            end
        end
    end
end
E_u = kron(sigeta,kron(sigeta,sigeta))*E_eps3;

% The auxiliary system z=[xf xs kron(xf,xf) xrd kron(xf,xs) kron(xf,kron(xf,xf))]
% The loading matrix
D = [gx+3/6*gssx*sig^2 gx 0.5*reshape(gxx,ny,nx^2) gx 2*0.5*reshape(gxx,ny,nx^2) 1/6*reshape(gxxx,ny,nx^3)];

numSim   = size(xSim,2);
conVar_y = zeros(ny,ny,numSim);
conStdY  = zeros(ny,numSim);
idx      = 0;
for s=1:numSim
    idx = idx + 1;
    if idx == 1000
        disp(['Draws so far: ', num2str(s)])
        idx = 0;
    end
    xf  = xSim(1:nx,s);
    xs  = xSim(nx+1:2*nx,s);
    %xrd = xSim(2*nx+1:3*nx,s);
    xf_xf = xf*xf';
    xf_xs = xf*xs';
    xf_xf_xf = reshape(kron3(xf_xf,xf),nx,nx,nx);

    % Computing variance of the innovations
    Var_xs = zeros(nx,nx);
    Var_xfxf = zeros(nx^2,nx^2);
    Var_xs_xfxf = zeros(nx,nx^2);
    Var_inov = Get_VarInov_3rd(hx,hxx,hss,sigeta,sig,...
        vectorMom3,vectorMom4,vectorMom5,vectorMom6,xs,xf_xf,xf_xs,xf_xf_xf,E_u,...
        Var_xs,Var_xfxf,Var_xs_xfxf,nx,ne);
    
    conVar_y(:,:,s) = D*Var_inov*D';
    conStdY(:,s) = sqrt(diag(conVar_y(:,:,s)));
    if any(abs(imag(conStdY(:,s)))>0)
        'here'
    end
end

%% The output
out.conVar_y    = conVar_y;
out.conStdY     = conStdY;
out.std_conStdY = std(conStdY,0,2);
if addBondYieldsON == 1
    labelYields = {};
    for i=1:length(model.maturities)
        labelYields{i} = ['$r^{(',num2str(model.maturities(1,i)),')}_t$'];
    end
    out.labely  = [model.labely,labelYields];
else
    out.labely  = model.labely;
end

end

